import re
import json
import argparse

class AsmParser:
    def __init__(self, asm_file, code_file, inline_file, block_file, output_file, sensitive_functions, missing_file):
        self.asm_file = asm_file
        self.code_file = code_file
        self.inline_file = inline_file  # 新增的 inline_disassembly.txt 文件
        self.block_file = block_file    # 新增的 block_info.txt 文件
        self.output_file = output_file
        self.missing_file = missing_file
        self.sensitive_functions = sensitive_functions
        self.result = {"plt_sensitive_address": {}}
        self.asm_data = ""
        self.code_lines = []
        self.inline_data = ""  # 存储 inline_disassembly.txt 文件内容
        self.block_data = ""   # 存储 block_info.txt 文件内容
        self.mem_addresses = {}
        self.call_results = []
        self.inline_results = []
        self.missing_instructions = []
        self.block_info = {}   # 存储解析后的基本块信息

    def read_files(self):
        """读取所有必要的文件"""
        with open(self.asm_file, 'r', encoding='utf-8') as file:
            self.asm_data = file.read()

        with open(self.code_file, 'r', encoding='utf-8') as file:
            self.code_lines = file.readlines()

        with open(self.inline_file, 'r', encoding='utf-8') as file:
            self.inline_data = file.read()

        with open(self.block_file, 'r', encoding='utf-8') as file:
            self.block_data = file.read()

    def parse_block_info(self):
        """解析 block_info.txt 文件，提取基本块信息"""
        func_pattern = re.compile(r'// Function:\s*(.*)')
        block_pattern = re.compile(r'Block Start: (0x[0-9a-fA-F]+), Block End: (0x[0-9a-fA-F]+)')

        current_func = ''
        blocks = []

        lines = self.block_data.splitlines()
        for line in lines:
            func_match = func_pattern.match(line)
            if func_match:
                # 保存之前函数的基本块信息
                if current_func:
                    self.block_info[current_func] = blocks
                current_func = func_match.group(1).strip()
                blocks = []
            else:
                block_match = block_pattern.match(line)
                if block_match:
                    block_start = block_match.group(1)
                    block_end = block_match.group(2)
                    blocks.append({"Block Start": block_start, "Block End": block_end})
        # 保存最后一个函数的基本块信息
        if current_func:
            self.block_info[current_func] = blocks

    def parse_data(self):
        """解析 ASM 文件内容，提取敏感函数的地址"""
        pattern = r'function name:([^\n]+)\n(0x[0-9A-Fa-f]+):\s*(.*?)\n'
        matches = re.findall(pattern, self.asm_data, re.MULTILINE)

        for match in matches:
            func_name = match[0]
            if func_name in self.sensitive_functions:
                mem_address = match[1]
                instruction = match[2]
                str_address = self.extract_address(instruction)
                self.result["plt_sensitive_address"][func_name] = {
                    "mem": mem_address,
                    "str": str_address
                }
                self.mem_addresses[mem_address.lower()] = func_name

    def extract_address(self, instruction):
        """从指令中提取地址信息"""
        str_address_match = re.search(r'\[(0x[0-9A-Fa-f]+)\]', instruction)
        if str_address_match:
            return str_address_match.group(1)
        else:
            str_address_match = re.search(r'0x[0-9A-Fa-f]+', instruction)
            if str_address_match:
                return str_address_match.group(0)
            else:
                return ""

    def find_calls(self):
        """在 ASM 文件中查找调用敏感函数的 call 指令"""
        lines = self.asm_data.splitlines()
        func_name_pattern = re.compile(r'^function name:(.+)')
        call_pattern = re.compile(r'^(0x[0-9A-Fa-f]+):\s*(call|jmp)\s+(0x[0-9A-Fa-f]+).*')
        current_func = ''
        existing_calls = set()

        for line in lines:
            func_match = func_name_pattern.match(line)
            if func_match:
                current_func = func_match.group(1)
                current_func = "sub_" + current_func[4:].upper() #SSH 修改为大写
            else:
                call_match = call_pattern.match(line)
                if call_match:
                    call_address = call_match.group(1)
                    instruction_type = call_match.group(2)
                    target_address = call_match.group(3)
                    target_address_lower = target_address.lower()
                    call_key = (call_address, target_address_lower)
                    if target_address_lower in self.mem_addresses and call_key not in existing_calls:
                        func_name = self.mem_addresses[target_address_lower]
                        record = {
                            "name": func_name,
                            "call_address": call_address,
                            "caller_func_name": current_func,
                            "inline": "no",
                            "related_par_address": {}
                        }
                        self.verify_and_find_params(record, target_address_lower)
                        self.call_results.append(record)
                        existing_calls.add(call_key)

    def verify_and_find_params(self, record, asm_target_address):
        """验证指令并查找相关参数地址，添加基本块信息"""
        call_address = record['call_address']
        asm_instruction = self.get_instruction_from_asm(call_address)
        code_instruction = self.get_instruction_from_code(call_address)

        asm_matches = re.search(r'call\s+(0x[0-9A-Fa-f]+)', asm_instruction, re.IGNORECASE) if asm_instruction else None
        code_matches = re.search(r'call\s+(\S+)', code_instruction, re.IGNORECASE) if code_instruction else None

        asm_calls_func = False

        if asm_matches:
            asm_called_address = asm_matches.group(1).lower()
            asm_calls_func = asm_called_address == asm_target_address

        code_calls_func = False

        if code_matches:
            code_called_name = code_matches.group(1).lstrip('_')
            code_calls_func = code_called_name == record['name']

        if not (asm_calls_func and code_calls_func):
            self.missing_instructions.append({
                "call_address": call_address,
                "asm_instruction": asm_instruction,
                "code_instruction": code_instruction,
                "issue": "Instructions do not match or are incorrect calls"
            })

        # 添加 caller_basic_block_address 信息
        caller_block = self.find_block_by_address(call_address, record['caller_func_name'])
        record["caller_basic_block_address"] = caller_block

        # 添加 caller_func 中的 basic_block 信息
        func_blocks = self.block_info.get(record['caller_func_name'], [])
        record["caller_func"] = {
            "caller_func_name": record['caller_func_name'],
            "basic_block": func_blocks
        }
        
        sensitive_functions1 = [        
        'strcpy', 'memcpy', 'memmove',
        'strcat', 'sprintf', 'gets',
        'wcscpy', 'wcscat' , 
        ]
        #sensitive_functions1 是参数个数为两个，参数注解为 dest 和 src 的敏感函数
        
        sensitive_functions2 = [        
        'memset', 
        ]
        #sensitive_functions2 是参数个数为三个，参数注解为 s , c 和 n 的敏感函数
        
        sensitive_functions3 = [        
        'fgets', 
        ]
        #sensitive_functions3 是参数个数为三个，参数注解为 s , n 和 stream 的敏感函数
        
        # sensitive_functions4 = [        
        # 'puts', 
        # ]
        #sensitive_functions4 是参数个数为一个，参数注解为 s 的敏感函数
        
        sensitive_functions5 = [        
        'recv', 
        ]
        #sensitive_functions5 是参数个数为四个，参数注解为 fd , n , flags 和 buf 的敏感函数
        
        sensitive_functions6 = [        
        'strncat', 'strncpy',
        ]
        #sensitive_functions6 是参数个数为三个，参数注解为 dest , src 和 n 的敏感函数
        
        # print(record['name'])
        
        # 添加代码信息
        addresses = [call_address.lower()]
        
        if record['name'] in sensitive_functions1:
            dest_address = self.find_previous_instruction(call_address, 'dest')
            src_address = self.find_previous_instruction(call_address, 'src')

            record['related_par_address'] = {
                "dest": dest_address,
                "src": src_address
            }
            
            addresses.append(src_address.lower())
            
        if record['name'] in sensitive_functions2:
            s_address = self.find_previous_instruction(call_address, 's')
            c_address = self.find_previous_instruction(call_address, 'c')
            n_address = self.find_previous_instruction(call_address, 'n')

            record['related_par_address'] = {
                "s": s_address,
                "c": c_address,
                "n": n_address
            }
            
            addresses.append(n_address.lower())
        
        if record['name'] in sensitive_functions3:
            s_address = self.find_previous_instruction(call_address, 's')
            n_address = self.find_previous_instruction(call_address, 'n')
            stream_address = self.find_previous_instruction(call_address, 'stream')

            record['related_par_address'] = {
                "s": s_address,
                "n": n_address,
                "stream": stream_address
            }
            
            addresses.append(stream_address.lower())
        
        # if record['name'] in sensitive_functions4:
        #     s_address = self.find_previous_instruction(call_address, 's')

        #     record['related_par_address'] = {  
        #         "s": s_address
        #     }
            
        #     addresses.append(s_address.lower())
            
        if record['name'] in sensitive_functions5:
            fd_address = self.find_previous_instruction(call_address, 'fd')
            n_address = self.find_previous_instruction(call_address, 'n')
            flags_address = self.find_previous_instruction(call_address, 'flags')
            buf_address = self.find_previous_instruction(call_address, 'buf')

            record['related_par_address'] = {
                "fd": fd_address,
                "n": n_address,
                "flags": flags_address,
                "buf": buf_address
            }
            
            addresses.append(buf_address.lower())
            
        if record['name'] in sensitive_functions6:
            dest_address = self.find_previous_instruction(call_address, 'dest')
            src_address = self.find_previous_instruction(call_address, 'src')
            n_address = self.find_previous_instruction(call_address, 'n')

            record['related_par_address'] = {
                "dest": dest_address,
                "src": src_address,
                "n": n_address
            }
            
            addresses.append(n_address.lower())
            

        address_lines = self.get_address_line_info(addresses)
        if len(address_lines) >= 2:
            # 按照地址顺序排序
            sorted_addresses = sorted(addresses, key=lambda x: int(x, 16))
            start_address = sorted_addresses[0]
            end_address = sorted_addresses[-1]
            code_info = self.get_code_block_info(start_address, end_address)
            record['code'] = code_info

    def find_block_by_address(self, address, func_name):
        """通过地址和函数名查找对应的基本块"""
        address_int = int(address, 16)
        func_blocks = self.block_info.get(func_name, [])
        for block in func_blocks:
            start_int = int(block['Block Start'], 16)
            end_int = int(block['Block End'], 16)
            if start_int <= address_int <= end_int:
                return {
                    "Block Start": block['Block Start'],
                    "Block End": block['Block End']
                }
        return {}

    def get_instruction_from_asm(self, address):
        """从 ASM 数据中获取指定地址的指令"""
        pattern = rf'^{address}:\s*(.+)$'
        match = re.search(pattern, self.asm_data, re.MULTILINE)
        if match:
            return match.group(1).strip()
        else:
            return None

    def get_instruction_from_code(self, address):
        """从 assembly_code.txt 中获取指定地址的指令"""
        for line in self.code_lines:
            line = line.strip()
            if line.lower().startswith(address.lower() + ':'):
                parts = line.split(':', 1)
                if len(parts) == 2:
                    return parts[1].strip()
        return None

    def find_previous_instruction(self, address, keyword):
        """在 assembly_code.txt 中从指定地址向上查找包含指定关键字的指令"""
        index = None
        address_lower = address.lower()
        for i, line in enumerate(self.code_lines):
            if line.strip().lower().startswith(address_lower + ':'):
                index = i
                break
        if index is None:
            self.missing_instructions.append({
                "address": address,
                "issue": f"Address {address} not found in assembly_code.txt"
            })
            return ""

        for i in range(index - 1, -1, -1):
            line = self.code_lines[i].strip()
            if ':' in line:
                addr, instr = line.split(':', 1)
                instr = instr.strip()
                if keyword.lower() in instr.lower():
                    return addr.strip()
        return ""
    
    def get_address_line_info(self, addresses):
        """获取指定地址列表在 asm 文件中的行号和内容"""
        lines = self.asm_data.splitlines()
        address_lines = {}
        for idx, line in enumerate(lines):
            line = line.strip()
            match = re.match(r'^(0x[0-9A-Fa-f]+):', line)
            if match:
                address = match.group(1).lower()
                if address in addresses:
                    address_lines[address] = {
                        'line_number': idx + 1,  # 行号从 1 开始
                        'content': line
                    }
                    if len(address_lines) == len(addresses):
                        break
        return address_lines
    
    def get_code_block_info(self, start_address, end_address):
        """获取从 start_address 到 end_address 之间的代码块信息"""
        lines = self.asm_data.splitlines()
        start_idx = end_idx = None
        for idx, line in enumerate(lines):
            line = line.strip()
            match = re.match(r'^(0x[0-9A-Fa-f]+):', line)
            if match:
                address = match.group(1).lower()
                if address == start_address:
                    start_idx = idx
                if address == end_address:
                    end_idx = idx
                    break
        if start_idx is not None and end_idx is not None:
            code_lines = lines[start_idx:end_idx + 1]
            start_columns = [len(line) - len(line.lstrip()) + 1 for line in code_lines]
            end_columns = [len(line.rstrip()) for line in code_lines]
            start_column = min(start_columns)
            end_column = max(end_columns)
            code_content = '\n'.join(code_lines)
            return {
                'start_line': start_idx + 1,
                'end_line': end_idx + 1,
                'start_column': start_column,
                'end_column': end_column,
                'code_content': code_content
            }
        else:
            return {}
        
    def get_function_code_info(self, func_name):
        """获取指定函数的代码块信息"""
        lines = self.asm_data.splitlines()
        start_idx = end_idx = None
        in_function = False
        for idx, line in enumerate(lines):
            line = line.strip()
            func_match = re.match(r'^function name:(.+)', line)
            if func_match:
                current_func = func_match.group(1).strip()
                if current_func == func_name:
                    in_function = True
                    start_idx = idx + 1  # 函数名下一行开始
                elif in_function:
                    # 已经找到目标函数，遇到下一个函数，结束
                    end_idx = idx - 1
                    break
            elif in_function and idx == len(lines) - 1:
                # 到达文件末尾
                end_idx = idx
        if start_idx is not None and end_idx is not None:
            code_lines = lines[start_idx:end_idx + 1]
            start_columns = [len(line) - len(line.lstrip()) + 1 for line in code_lines if line.strip()]
            end_columns = [len(line.rstrip()) for line in code_lines if line.strip()]
            if start_columns and end_columns:
                start_column = min(start_columns)
                end_column = max(end_columns)
            else:
                start_column = end_column = 1
            code_content = '\n'.join(code_lines)
            return {
                'start_line': start_idx + 1,
                'end_line': end_idx + 1,
                'start_column': start_column,
                'end_column': end_column,
                'code_content': code_content
            }
        else:
            return {}

    def parse_inline_disassembly(self):
        """解析 inline_disassembly.txt 文件，提取 INLINED 部分的信息"""
        # pattern = r"// Function '(\w+)' in '([\w\d_]+)' - INLINED:\n(?://.*\n)*// Near address: 0x([0-9A-Fa-f]+)"
        pattern = (
            r"// Function '(\w+)' in '([\w\d_]+)' - INLINED:\n"
            r"// Occurrence in decompiled code at offset \d+\n"
            r"Inline Decompile Code '(.+?)'\n"
            r"'0x([0-9A-Fa-f]+): .*'\n"  # 修改此行以捕获主要地址
            r"// Near address: 0x([0-9A-Fa-f]+)\n"
            r"((?:0x[0-9A-Fa-f]+: .+\n)*)"
        )  
        matches = re.finditer(pattern, self.inline_data, re.MULTILINE)

        match_count = 0  # 调试用
        for match in matches:
            match_count += 1
            sensitive_func = match.group(1)
            # print(sensitive_func)
            caller_func = match.group(2)
            # print(caller_func)
            decompile_code = match.group(3)
            # print(decompile_code)
            address = match.group(4)
            # print(address)
            additional_code = match.group(5) if match.group(5) else ""
            # print(f"匹配到的函数: {sensitive_func}，调用者: {caller_func}")  # 调试用
            # print(f"反编译代码: {decompile_code}")  # 调试用
            # print(f"地址: {address}")  # 调试用
            # print(f"附加代码: {additional_code}")  # 调试用

            if sensitive_func in self.sensitive_functions:
                # 添加 basic_block 信息
                func_blocks = self.block_info.get(caller_func, [])
                # 获取函数代码信息
                code_info = self.get_function_code_info(caller_func)
                record = {
                    "inline": "yes",
                    "inline_func_name_predict": sensitive_func,
                    "call_address": f"0x{address}",
                    "caller_func": {
                        "name": caller_func,
                        "basic_block": func_blocks
                    },
                    "code": code_info
                }
                self.inline_results.append(record)
        print(f"匹配到的 INLINED 项目数量: {match_count}")  # 调试用

    def update_result(self):
        """更新 result 字典，添加 direct_sensitive_func 和 inline_sensitive_func 信息"""
        self.result["direct_sensitive_func"] = self.call_results
        self.result["inline_sensitive_func"] = self.inline_results

    def save_result(self):
        """将结果保存到 JSON 文件"""
        with open(self.output_file, 'w', encoding='utf-8') as f:
            json.dump(self.result, f, ensure_ascii=False, indent=4)
        with open(self.missing_file, 'w', encoding='utf-8') as f:
            json.dump(self.missing_instructions, f, ensure_ascii=False, indent=4)
        print(f"结果已保存到文件 {self.output_file}")
        print(f"缺失信息已保存到文件 {self.missing_file}")

    def run(self):
        """运行解析器"""
        self.read_files()
        self.parse_block_info()  # 解析 block_info.txt
        self.parse_data()
        self.find_calls()
        self.parse_inline_disassembly()
        self.update_result()
        self.save_result()

# if __name__ == "__main__":
#     sensitive_functions = [
#         'strcpy', 'memcpy', 'memmove',
#         'strcat', 'sprintf', 'gets',
#         'wcscpy', 'wcscat'
#     ]
#     parser = AsmParser(
#         asm_file=r'D:\Users\Administrator\Desktop\study\project\dataset\note_test\json\bad1\CWE121_Stack_Based_Buffer_Overflow__CWE805_struct_declare_memcpy_82\CWE121_Stack_Based_Buffer_Overflow__CWE805_struct_declare_memcpy_82.asm',
#         code_file=r'D:\Users\Administrator\Desktop\study\project\dataset\note_test\json\bad1\CWE121_Stack_Based_Buffer_Overflow__CWE805_struct_declare_memcpy_82\CWE121_Stack_Based_Buffer_Overflow__CWE805_struct_declare_memcpy_82_assembly_code.txt',
#         inline_file=r'D:\Users\Administrator\Desktop\study\project\dataset\note_test\json\bad1\CWE121_Stack_Based_Buffer_Overflow__CWE805_struct_declare_memcpy_82\CWE121_Stack_Based_Buffer_Overflow__CWE805_struct_declare_memcpy_82_inline_disassembly.txt',
#         block_file=r'D:\Users\Administrator\Desktop\study\project\dataset\note_test\json\bad1\CWE121_Stack_Based_Buffer_Overflow__CWE805_struct_declare_memcpy_82\CWE121_Stack_Based_Buffer_Overflow__CWE805_struct_declare_memcpy_82_block_info.txt',   # 添加 block_info.txt 文件路径
#         output_file=r'D:\Users\Administrator\Desktop\study\project\dataset\note_test\json\bad1\CWE121_Stack_Based_Buffer_Overflow__CWE805_struct_declare_memcpy_82\result.json',
#         missing_file=r'D:\Users\Administrator\Desktop\study\project\dataset\note_test\json\bad1\CWE121_Stack_Based_Buffer_Overflow__CWE805_struct_declare_memcpy_82\address_mismatch.json',
#         sensitive_functions=sensitive_functions
#     )
#     parser.run()
    
def main():
    parser = argparse.ArgumentParser(description='ASM Parser Script')
    parser.add_argument('--asm_file', required=True, help='Path to the ASM file')
    parser.add_argument('--code_file', required=True, help='Path to the code file')
    parser.add_argument('--inline_file', required=True, help='Path to the inline disassembly file')
    parser.add_argument('--block_file', required=True, help='Path to the block info file')
    parser.add_argument('--output_file', required=True, help='Path to save the result JSON')
    parser.add_argument('--missing_file', required=True, help='Path to save the address mismatch JSON')

    args = parser.parse_args()

    sensitive_functions = [
        'strcpy', 'memcpy', 'memmove',
        'strcat', 'sprintf', 'gets',
        'wcscpy', 'wcscat' , 'memset' ,
        'fgets', 'recv',
        'strncat', 'strncpy',
    ]

    parser_instance = AsmParser(
        asm_file=args.asm_file,
        code_file=args.code_file,
        inline_file=args.inline_file,
        block_file=args.block_file,
        output_file=args.output_file,
        missing_file=args.missing_file,
        sensitive_functions=sensitive_functions
    )

    parser_instance.run()

if __name__ == "__main__":
    main()